# Generated from C:/workspace/mathsolver_engine\mathparser.g4 by ANTLR 4.6
from antlr4 import *
from sympy import sympify,Interval,S,Union,Set
from collections import defaultdict
from mathsolver.mathparser.mathparserVisitor import mathparserVisitor
# This class defines a complete generic visitor for a parse tree produced by mathparserParser.

class myparser(mathparserVisitor):
    expr_head = ["Plus","Rational","Times","Power"]
    mapping = {}

    # Visit a parse tree produced by mathparserParser#A1.
    def visitA1(self, ctx):
        return ctx.getText()


    # Visit a parse tree produced by mathparserParser#A3.
    def visitA2(self, ctx):
        return "-(%s)" % self.visit(ctx.atom())

    # Visit a parse tree produced by mathparserParser#T1.
    def visitT1(self, ctx):
        flag = ctx.getChild(0).getText()
        tokens = [self.visit(token) for token in ctx.token()]
        if "Or" in flag:
            values = defaultdict(list)
            for ans in tokens:
                if isinstance(ans, tuple):
                    key, value = ans
                else:
                    key, value = list(ans.items())[0]
                values[key].append(value)

            if not isinstance(key, tuple) and isinstance(values[key][0], Set):
                values[key] = Union(values[key])

            return dict(values)

        elif "And" in flag:
            tokens = dict(tokens)
            if len(tokens) == 1:
                return tokens
            else:
                syms = tokens.keys()
                values = tokens.values()
                return {tuple(syms): tuple(values)}


    def map(self, flag):
        if flag == "Equal" or flag == "Rule":
            return "=="
        elif flag == "Unequal":
            return "!="
        elif flag == "Less":
            return "<"
        elif flag == "LessEqual":
            return "<="
        elif flag == "Greater":
            return ">"
        elif flag == "GreaterEqual":
            return ">="
        else:
            return flag

    def visitT2(self, ctx):
        flag = self.map(ctx.getChild(0).getText())
        tokens = [self.visit(token) for token in ctx.token()]
        return self._visitT2(flag,tokens)

    def _visitT2(self,flag, tokens):
        if flag == "Inequality":
            sym = sympify(tokens[2]).subs(self.mapping)
            return (sym,Interval(sympify(tokens[0]).subs(self.mapping), sympify(tokens[4]).subs(self.mapping), "=" not in tokens[1], "=" not in tokens[3]))
        else:
            sym = sympify(tokens[0]).subs(self.mapping)
            value = sympify(tokens[1]).subs(self.mapping)

            if flag == "==":
                return (sym, value)
            elif ">" in flag:
                return (sym, Interval(value, S.Infinity, "=" not in flag))
            elif "<" in flag:
                return (sym, Interval(-S.Infinity, value, True, "=" not in flag))
            else:
                return

    # Visit a parse tree produced by mathparserParser#T3.
    def visitT3(self, ctx):
        flag = ctx.getChild(0).getText()
        tokens = [self.visit(token) for token in ctx.token()]
        if flag == "Log":
            return "log(%s)" % tokens[0]
        elif flag == "Abs":
            return "Abs(%s)"  % tokens[0]
        elif flag == "Subscript":
            return tokens[0] + "_(" + tokens[1] + ")"
        elif flag == "DirectedInfinity":
            return tokens(0) + "*oo"
        elif flag == "Slot":
            return "x"
        elif flag == "Function":
            return tokens[0]
        elif flag == "Root":
            return "Root(" + ",".join(tokens) + ")"
        elif flag in self.expr_head:
            if flag == "Plus":
                symbol = "+"
            elif flag == "Rational":
                symbol = "/"
            elif flag == "Times":
                symbol = "*"
            elif flag == "Power":
                symbol = "**"
            else:
                symbol = ""

            if symbol == "*" and tokens[0] == "-1":
                return "-(" + symbol.join(tokens[1:]) + ")"
            else:
                return "(" + symbol.join(tokens) + ")"
        else:
            return ctx.getText()


    # Visit a parse tree produced by mathparserParser#T4.
    def visitT4(self, ctx):
        return self.visitChildren(ctx)

    def visitT5(self, ctx):
        return self.map(ctx.getText())


if __name__ == '__main__':
    from antlr4 import *
    from mathsolver.mathparser.mathparserLexer import mathparserLexer
    from mathsolver.mathparser.mathparserParser import mathparserParser
    lexer = mathparserLexer(InputStream("Or[LessEqual[a, 0], And[Greater[a, 0], LessEqual[x, Times[-1, Power[3, Rational[-1, 2]], Power[a, Rational[1, 2]]]]], And[Greater[a, 0], GreaterEqual[x, Times[Power[3, Rational[-1, 2]], Power[a, Rational[1, 2]]]]]]"))
    stream = CommonTokenStream(lexer)
    parser = mathparserParser(stream)
    tree = parser.token()
    # print myparser().visit(tree)